package com.lucifer.hawkeye.ai.service.impl;

import com.lucifer.hawkeye.ai.advisor.HawkeyeVectorStoreChatMemoryAdvisor;
import com.lucifer.hawkeye.ai.config.HawkeyeAiConfig;
import com.lucifer.hawkeye.ai.domain.conversation.bo.ChatMemoryBo;
import com.lucifer.hawkeye.ai.domain.nl2sql.type.TrainType;
import com.lucifer.hawkeye.ai.model.ChatClientFactory;
import com.lucifer.hawkeye.ai.nl2sql.HawkeyeSqlEngine;
import com.lucifer.hawkeye.ai.prompt.HawkeyePrompt;
import com.lucifer.hawkeye.ai.rag.HawkeyeRagEngine;
import com.lucifer.hawkeye.ai.rag.RagEngine;
import com.lucifer.hawkeye.ai.service.ExecuteSqlService;
import com.lucifer.hawkeye.ai.service.Nl2SqlService;
import com.lucifer.hawkeye.ai.vector.HawkeyeVectorStoreEngine;
import com.lucifer.hawkeye.ai.vector.VectorStoreEngine;
import cn.hutool.core.util.StrUtil;
import com.alibaba.fastjson2.JSONObject;
import com.baomidou.dynamic.datasource.DynamicRoutingDataSource;
import com.google.common.collect.Lists;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.SystemPromptTemplate;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.id.RandomIdGenerator;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.jdbc.core.BeanPropertyRowMapper;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.stereotype.Service;

import java.util.List;
import java.util.Map;
import java.util.Optional;

/**
 * @author lucifer
 * @date 2025/4/1 16:57
 */
@Slf4j
@Service
public class Nl2SqlServiceImpl implements Nl2SqlService {

    @Resource
    private HawkeyeAiConfig hawkeyeAiConfig;


    @Resource(name = "nl2sqlVectorStore")
    private VectorStore nl2sqlVectorStore;


    @Resource(name = "chatMemoryVectorStore")
    private VectorStore chatMemoryVectorStore;


    @Resource
    private DynamicRoutingDataSource dynamicRoutingDataSource;

    @Resource
    private ExecuteSqlService executeSqlService;



    @Override
    public Boolean train(String conversationId, String trainType, String question, String content,String datasource, String tableName) {
        if (!TrainType.check(trainType)) {
            throw new IllegalArgumentException("Invalid train type: " + trainType);
        }
        if(!dynamicRoutingDataSource.getDataSources().containsKey(datasource)) {
            throw new IllegalArgumentException("Invalid datasource: " + datasource);
        }
        String trainTypeCode = TrainType.codeOf(trainType).getCode();
        Document document = null;
        Map<String, Object> metadata = Map.of(
                "conversationId", conversationId,
                "rag_type", trainType,
                "semantic", question,
                "datasource", datasource,
                "table_name", tableName);
        if (TrainType.DDL.getCode().equals(trainTypeCode)) {
            JdbcTemplate jdbcTemplate = (JdbcTemplate)nl2sqlVectorStore.getNativeClient().orElseThrow(() -> new IllegalArgumentException("Nl2SqlVectorStore is null"));
            List<String> ids = jdbcTemplate.queryForList("SELECT id  FROM nl2sql_store WHERE metadata->>'conversationId' = ? AND metadata->>'rag_type' = 'DDL'", String.class, conversationId);
            document = new Document(ids.isEmpty()? new RandomIdGenerator().generateId(): ids.get(0), content,metadata);
        }
        if (TrainType.SQL.getCode().equals(trainTypeCode)) {
            JSONObject object = new JSONObject();
            object.put("question", question);
            object.put("sql", content);
            document = new Document(object.toJSONString(), metadata);
        }
        VectorStoreEngine storeEngine = HawkeyeVectorStoreEngine.builder().vectorStore(nl2sqlVectorStore).build();
        RagEngine ragEngine = HawkeyeRagEngine.builder().vectorStoreEngine(storeEngine).build();
        ragEngine.add(document);
        return true;

    }

    @Override
    public Map getTrainMetadata(String conversationId) {
        JdbcTemplate jdbcTemplate = (JdbcTemplate)nl2sqlVectorStore.getNativeClient().orElseThrow(() -> new IllegalArgumentException("Nl2SqlVectorStore is null"));
        BeanPropertyRowMapper<ChatMemoryBo> nl2SqlBoBeanPropertyRowMapper = new BeanPropertyRowMapper<>(ChatMemoryBo.class);
        ChatMemoryBo nl2SqlBo = jdbcTemplate.queryForObject("SELECT id, content, metadata AS metadataJson, embedding FROM nl2sql_store WHERE metadata->>'conversationId' = ? AND metadata->>'rag_type' = 'DDL'", nl2SqlBoBeanPropertyRowMapper, conversationId);
        nl2SqlBo = new ChatMemoryBo(nl2SqlBo);
        return nl2SqlBo.getMetadata();
    }

    @Override
    public String generateSql(String conversationId, String content,String datasource, String tableName) {
        VectorStoreEngine storeEngine = HawkeyeVectorStoreEngine.builder().vectorStore(nl2sqlVectorStore).build();
        RagEngine ragEngine = HawkeyeRagEngine.builder().vectorStoreEngine(storeEngine).build();
        HawkeyeSqlEngine sqlEngine = HawkeyeSqlEngine.builder().ragEngine(ragEngine).chatModel(hawkeyeAiConfig.chatModel()).executeSqlService(executeSqlService).build();
        String sql = sqlEngine.generateSql(conversationId, content,datasource, tableName);
        log.info("会话ID = {}，问题 = {}，生成的SQL = {}", conversationId,content,sql);
        List<Map<String, Object>> data = Lists.newArrayList();
        try {
            data = sqlEngine.executeSql(datasource,sql);
        } catch (Exception e) {
            log.error("会话ID = {}，问题 = {}，生成的SQL = {}，执行SQL语句异常，异常原因 = {}", conversationId,content,sql,e.getMessage());
        }
        String summary  = generateSummary(conversationId, content, sql, data.toString());
        return summary;
    }

    @Override
    public String generateSummary(String conversationId, String question, String sql, String context) {
        SystemPromptTemplate summarySystemPromptTemplate = HawkeyePrompt.generateSummarySystemPromptTemplate();
        Message systemMessage = summarySystemPromptTemplate.createMessage(Map.of(
                "question", getOrDefault(question, "No question provided"),
                "context", getOrDefault(context, "No data available")
        ));
        UserMessage userMessage = new UserMessage(question);
        HawkeyeVectorStoreChatMemoryAdvisor storeChatMemoryAdvisor = HawkeyeVectorStoreChatMemoryAdvisor.builder(chatMemoryVectorStore).build();
        UserMessage beforeUserMessage = new UserMessage(question);
        List<Document> beforeDocuments = storeChatMemoryAdvisor.toDocuments(List.of(beforeUserMessage), conversationId);
        chatMemoryVectorStore.write(beforeDocuments);
        Prompt prompt = new Prompt(List.of(systemMessage, userMessage));
        String text = "";
        try {
            ChatClient.CallResponseSpec call = ChatClientFactory.buildChatClient(hawkeyeAiConfig.chatModel()).prompt(prompt).call();
            text = call.chatResponse().getResult().getOutput().getText();
        } catch (Exception e) {
            log.error("会话ID = {}，问题 = {}，生成的SQL = {}，数据分析异常，异常原因 = {}", conversationId,question,sql,e.getMessage());
        }
        String format = outSummary(question,sql,text);
        AssistantMessage afterAssistantMessage = new AssistantMessage(format);
        List<Document> afterDocuments = storeChatMemoryAdvisor.toDocuments(List.of(afterAssistantMessage), conversationId);
        chatMemoryVectorStore.write(afterDocuments);
        return format;
    }

    private static String getOrDefault(String value, String defaultValue) {
        return Optional.ofNullable(value).orElse(defaultValue);
    }

    private static String outSummary(String... params) {
        String res = """
            #### 原始问题:
            ```tex
            {}
            ```
            #### 生成SQL:
            ```sql
            {}
            ```
            #### 执行SQL结果总结:
         
            {}
           
            """;
        String format = StrUtil.format(res, params[0],  params[1], params[2]);
        return format;
    }
}
